import torch

"""
只支持正方形的图片输入。
学习率=0.05左右
"""


class AlexNet(torch.nn.Module):
    def __init__(self, size, channels=3):
        super(AlexNet, self).__init__()
        self._size = size
        self._channels = channels
        self._dropout = torch.nn.Dropout(p=0.5)
        self._active = torch.nn.ReLU()
        self._pooling = torch.nn.MaxPool2d(kernel_size=3, stride=2)
        self._flatten = torch.nn.Flatten()
        self._conv1 = torch.nn.Conv2d(channels, 96, kernel_size=(11,), stride=(4,), padding=1)
        self._conv2 = torch.nn.Conv2d(96, 256, kernel_size=(5,), padding=2)
        self._conv3 = torch.nn.Conv2d(256, 384, kernel_size=(3,), padding=1)
        self._conv4 = torch.nn.Conv2d(384, 384, kernel_size=(3,), padding=1)
        self._conv5 = torch.nn.Conv2d(384, 384, kernel_size=(3,), padding=1)
        self._linear1 = torch.nn.Linear((((((size - 8) // 4 - 1) // 2 - 1) // 2 - 1) // 2) ** 2 * 384, 4096)  # 自动计算size，正确性有待验证。
        self._linear2 = torch.nn.Linear(4096, 4096)
        self._linear3 = torch.nn.Linear(4096, 10)

    def forward(self, x):
        x = x.view(-1, self._channels, self._size, self._size)
        x = self._conv1(x)
        x = self._active(x)
        x = self._pooling(x)
        x = self._conv2(x)
        x = self._active(x)
        x = self._pooling(x)
        x = self._conv3(x)
        x = self._active(x)
        x = self._conv4(x)
        x = self._active(x)
        x = self._conv5(x)
        x = self._active(x)
        x = self._pooling(x)
        x = self._linear1(x)
        x = self._active(x)
        x = self._dropout(x)
        x = self._linear2(x)
        x = self._active(x)
        x = self._dropout(x)
        return self._linear3(x)
